import tensorflow as tf
from tensorflow import keras
from PIL import Image, ImageDraw
from ultralytics import YOLO
import numpy as np
import relic
from utils import image_util

boxer = YOLO("boxer.pt")
textRecognizer = keras.models.load_model("textRecognitionKeras")
class_names = None
with open("class_names.txt", mode="r", encoding="utf-8") as f:
    class_names = f.readlines()

for i in range(len(class_names)):
    class_names[i] = class_names[i].replace("\n", "")

print(class_names)


# list[PIL Image]
def get_text_by_img(images):
    images_np = []
    # if isinstance(images, list):
    #     pass
    # elif isinstance(images, Image.Image):
    #     images = [images]

    for image in images:
        images_np.append(np.array(image_util.convert_image(image))/255)

    images_np = np.array(images_np)
    y_predict = textRecognizer(images_np,training=False)
    y_predict1 = []
    for i in y_predict:
        i = i.numpy()
        y_predict1.append(class_names[np.where(i == i.max())[0][0]])

    return y_predict1


def get_relic_by_image(image: Image.Image) -> relic.Relic:
    image = image.resize(size=(410, 460))
    result = boxer(image)
    # Image.fromarray(result[0].plot()).show()
    relic0 = relic.Relic()
    for i in result[0].boxes.xyxy[result[0].boxes.cls == list(result[0].names.values()).index("relic_type")]:
        relic0.relic_type = get_text_by_img([image.crop(i.tolist())])[0]

    for i in result[0].boxes.xyxy[result[0].boxes.cls == list(result[0].names.values()).index("main_entry")]:
        relic0.main_entry_type = get_text_by_img([image.crop(i.tolist())])[0]

    for i in result[0].boxes.xyxy[result[0].boxes.cls == list(result[0].names.values()).index("second_lemma")]:
        relic0.secondary_lemma_type.append(get_text_by_img([image.crop(i.tolist())])[0])

    for i in result[0].boxes.xyxy[result[0].boxes.cls == list(result[0].names.values()).index("coordinate_name")]:
        relic0.coordinate_name = get_text_by_img([image.crop(i.tolist())])[0]

    relic0.locked = (5 in result[0].boxes.cls)
    return relic0


if __name__ == '__main__':
    img = Image.open("datasets/relic_images/images/train/0.png")
    relic1 = get_relic_by_image(img)
    print(relic1.tostring1())
    print("aaa")
